package com.ocean.everyday.d002_findingmkaverage;

import com.ocean.linkedlist.ListNode;

import java.util.Deque;
import java.util.LinkedList;
import java.util.TreeMap;

/**
 * https://leetcode.cn/problems/finding-mk-average/
 * 求出MK平均值
 */
public class MKAverage {
    private int m;
    private int k;
    // 这里 k个最小的
    private TreeMap<Integer, Integer> s1;
    private int size1;
    // 这里保存m−2k 个元素
    private TreeMap<Integer, Integer> s2;
    private int size2;
    // 这里保存k个最大的
    private TreeMap<Integer, Integer> s3;
    private int size3;
    // 这里只会保留m个元素 多了就删除
    private Deque<Integer> deque;
    private long sum2;

    public static void main(String[] args) {
        MKAverage obj = new MKAverage(3, 1);
        obj.addElement(3);        // 当前元素为 [3]
        obj.addElement(1);        // 当前元素为 [3,1]
        int i = obj.calculateMKAverage();// 返回 -1 ，因为 m = 3
        System.out.println(i);
        obj.addElement(10);
        i = obj.calculateMKAverage();
        System.out.println(i);

        obj.addElement(5);        // 当前元素为 [3,1,10,5]
        obj.addElement(5);        // 当前元素为 [3,1,10,5,5]
        obj.addElement(5);        // 当前元素为 [3,1,10,5,5,5]
        i = obj.calculateMKAverage();
        System.out.println(i);
    }


    public MKAverage(int m, int k) {
        this.m = m;
        this.k = k;
        this.size1 = 0;
        this.size2 = 0;
        this.size3 = 0;
        this.s1 = new TreeMap<>();
        this.s2 = new TreeMap<>();
        this.s3 = new TreeMap<>();
        deque = new LinkedList<>();
    }

    public void addElement(int num) {
        deque.offer(num);
        if (deque.size() <= m) {
            s2.put(num, s2.getOrDefault(num, 0) + 1);
            sum2 += num;
            size2++;
            // 如果元素有m个 就要分到不同的集合中
            if (deque.size() == m) {
                // 拿出k个到s1中
                while (size1 < k) {
                    Integer cur = s2.firstKey();
                    s1.put(cur, s1.getOrDefault(cur, 0) + 1);
                    size1++;
                    s2.put(cur, s2.get(cur) - 1);
                    sum2 -= cur;
                    if (s2.get(cur) == 0) {
                        s2.remove(cur);
                    }
                }
                while (size3 < k) {
                    Integer cur = s2.lastKey();
                    s3.put(cur, s3.getOrDefault(cur, 0) + 1);
                    size3++;
                    s2.put(cur, s2.get(cur) - 1);
                    sum2 -= cur;
                    if (s2.get(cur) == 0) {
                        s2.remove(cur);
                    }
                }
            }
            return;
        }
        // m+1个后:先判断当前元素加入到哪一边
        // 比最小集合的最大值还小
        if (num < s1.lastKey()) {
            // 说明这个元素要加入到s1中 要将s1的最后一个元素放入到deque中
            Integer last = s1.lastKey();
            s1.put(num, s1.getOrDefault(num, 0) + 1);
            s1.put(last, s1.get(last) - 1);
            if (s1.getOrDefault(last, 0) == 0) {
                s1.remove(last);
            }
            s2.put(last, s2.getOrDefault(last, 0) + 1);
            sum2 += last;
            size2++;
        } else if (num > s3.firstKey()) { // 比最大集合的最小值大 要放到最大集合中
            Integer firstKey = s3.firstKey();
            s3.put(num, s3.getOrDefault(num, 0) + 1);
            s3.put(firstKey, s3.get(firstKey) - 1);
            if (s3.getOrDefault(firstKey, 0) == 0) {
                s3.remove(firstKey);
            }
            s2.put(firstKey, s2.getOrDefault(firstKey, 0) + 1);
            sum2 += firstKey;
            size2++;
        } else {
            // 直接放到s2中
            s2.put(num, s2.getOrDefault(num, 0) + 1);
            sum2 += num;
            size2++;
        }
        // 这个时候s2中就会多出一个元素 需要移除
        Integer poll = deque.poll();
        // 要判断拉出来的这个元素是在哪个集合中
        if (s1.containsKey(poll)) {
            s1.put(poll, s1.get(poll) - 1);
            if (s1.getOrDefault(poll, 0) == 0) {
                s1.remove(poll);
            }
            // 将s2中的第一个放到s1
            Integer first = s2.firstKey();
            s1.put(first, s1.getOrDefault(first, 0) + 1);
            s2.put(first, s2.get(first) - 1);
            if (s2.getOrDefault(first, 0) == 0) {
                s2.remove(first);
            }
            size2--;
            sum2 -= first;
        } else if (s3.containsKey(poll)) {
            s3.put(poll, s3.get(poll) - 1);
            if (s3.getOrDefault(poll, 0) == 0) {
                s3.remove(poll);
            }
            // 将s2中的最后一个放到s3
            Integer last = s2.lastKey();
            s3.put(last, s3.getOrDefault(last, 0) + 1);
            s2.put(last, s2.get(last) - 1);
            if (s2.getOrDefault(last, 0) == 0) {
                s2.remove(last);
            }
            size2--;
            sum2 -= last;
        } else {
            // 直接从s2移除
            s2.put(poll, s2.get(poll) - 1);
            if (s2.get(poll) == 0) {
                s2.remove(poll);
            }
            sum2 -= poll;
            size2--;
        }

    }

    /**
     * @return
     */
    public int calculateMKAverage() {
        if (deque.size() < m) {
            return -1;
        }
        return (int) sum2 / (m - 2 * k);
    }
}
